import random
import time
import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
from gurobipy import Model, GRB, LinExpr
import os

# ========== STEP 1: GIFT Encryption and Plaintext-Ciphertext Pair Generation ==========

def gift_encrypt(plaintext, key, rounds):
    """Simplified GIFT encryption (substitute with actual implementation)."""
    ciphertext = plaintext ^ key  # Simplified encryption (replace with actual GIFT rounds)
    return ciphertext

def generate_plaintext_ciphertext_pairs(n_pairs, key, rounds):
    """Generate n_pairs of plaintext-ciphertext pairs."""
    pairs = []
    plaintexts = []
    for _ in range(n_pairs):
        plaintext = random.randint(0, 0xFFFF)  # Generate a random 16-bit plaintext
        ciphertext = gift_encrypt(plaintext, key, rounds)  # Encrypt to get ciphertext
        pairs.append((plaintext, ciphertext))
        plaintexts.append(plaintext)  # Collect plaintexts for visualization
    return pairs, plaintexts

# ========== STEP 2: Calculate Differences and Differential Distribution ==========

def calculate_differences(pairs):
    """Calculate differences between plaintexts and ciphertexts."""
    plaintext_diffs = []
    ciphertext_diffs = []
    for (p1, c1), (p2, c2) in zip(pairs[:-1], pairs[1:]):
        # Calculate plaintext difference (XOR of two plaintexts)
        plaintext_diffs.append(p1 ^ p2)
        # Calculate ciphertext difference (XOR of two ciphertexts)
        ciphertext_diffs.append(c1 ^ c2)
    
    return Counter(plaintext_diffs), Counter(ciphertext_diffs)

# Generate Diffusion Distribution Table (DDT) for the S-box
def generate_ddt(sbox):
    size = len(sbox)
    ddt = np.zeros((size, size), dtype=int)
    for x1 in range(size):
        for x2 in range(size):
            diff_in = x1 ^ x2
            diff_out = sbox[x1] ^ sbox[x2]
            ddt[diff_in][diff_out] += 1
    return ddt

# ========== STEP 3: Traditional Attack ==========

def traditional_attack(pairs, ddt, target_diff):
    """Traditional differential attack simulation."""
    potential_keys = []
    for candidate_key in range(0xFFFF):  # Example: Brute force all possible keys
        matches = 0
        for plaintext, ciphertext in pairs:
            decrypted = ciphertext ^ candidate_key  # Simple decryption step
            # If the decrypted value matches the target difference
            if decrypted == plaintext ^ target_diff:
                matches += 1
        if matches > len(pairs) * 0.7:  # Example threshold (e.g., 70% matches)
            potential_keys.append(candidate_key)
    return potential_keys

# ========== STEP 4: MILP Optimization and Improved Attack ==========

def milp_optimized_attack(pairs, ddt, target_diff, sbox_size=16):
    """Simulate MILP optimization using Gurobi."""
    model = Model("milp_attack")

    # Variables for each possible difference (binary decision variable for each pair of differences)
    vars = {}
    for i in range(sbox_size):
        for j in range(sbox_size):
            vars[(i, j)] = model.addVar(vtype=GRB.BINARY, name=f"x_{i}_{j}")

    # Constraints: Each difference in the DDT should be considered only once in the optimal path.
    for i in range(sbox_size):
        for j in range(sbox_size):
            model.addConstr(vars[(i, j)] <= ddt[i][j])

    # Objective function: Maximize the sum of the selected paths' probabilities.
    obj = LinExpr()
    for i in range(sbox_size):
        for j in range(sbox_size):
            obj += vars[(i, j)] * ddt[i][j]
    
    model.setObjective(obj, GRB.MAXIMIZE)

    # Solve the model
    model.optimize()

    # Extract the results
    selected_diffs = []
    if model.status == GRB.OPTIMAL:
        for i in range(sbox_size):
            for j in range(sbox_size):
                if vars[(i, j)].x > 0.5:  # if the difference path is selected
                    selected_diffs.append((i, j))

    return selected_diffs

# ========== STEP 5: Statistical Analysis and Visualization ==========

def calculate_path_probability(path, ddt):
    """Calculate the probability of a differential path."""
    prob = 1.0
    for (delta_in, delta_out) in path:
        prob *= ddt[delta_in][delta_out] / sum(ddt[delta_in])
    return prob

def save_to_file(data, filename):
    """Save data to a text file."""
    with open(filename, 'w') as file:
        for line in data:
            file.write("\t".join(map(str, line)) + "\n")

def visualize_results_and_save(success_rates, probabilities, times, plaintexts, pairs, pt_diffs, ct_diffs, output_dir):
    """Visualize and save all the results."""
    labels = ["Traditional Attack", "Improved Attack"]
    
    # Success Rates Plot
    plt.bar(labels, success_rates, color=["blue", "green"])
    plt.title("Success Rate Comparison")
    plt.ylabel("Success Rate (%)")
    plt.savefig(f"{output_dir}/success_rate_comparison.png")
    plt.clf()  # Clear the figure to avoid overlap
    
    # Path Probabilities Plot
    plt.bar(labels, probabilities, color=["orange", "purple"])
    plt.title("Path Probability Comparison")
    plt.ylabel("Probability")
    plt.savefig(f"{output_dir}/path_probability_comparison.png")
    plt.clf()
    
    # Efficiency (Time) Plot
    plt.bar(labels, times, color=["red", "cyan"])
    plt.title("Efficiency Comparison (Time)")
    plt.ylabel("Average Time (s)")
    plt.savefig(f"{output_dir}/efficiency_comparison.png")
    plt.clf()
    
    # Visualize plaintexts (random data)
    plt.hist(plaintexts, bins=50, color='blue', alpha=0.7)
    plt.title("Generated Plaintext Distribution")
    plt.xlabel("Plaintext Value")
    plt.ylabel("Frequency")
    plt.savefig(f"{output_dir}/plaintext_distribution.png")
    plt.clf()
    
    # Visualize differences
    plt.figure(figsize=(8, 6))
    plt.bar(pt_diffs.keys(), pt_diffs.values(), color="blue", alpha=0.7, label="Plaintext Differences")
    plt.bar(ct_diffs.keys(), ct_diffs.values(), color="orange", alpha=0.7, label="Ciphertext Differences")
    plt.title("Plaintext and Ciphertext Difference Distributions")
    plt.xlabel("Difference (XOR value)")
    plt.ylabel("Frequency")
    plt.legend()
    plt.savefig(f"{output_dir}/difference_distributions.png")
    plt.clf()
    
    print(f"Images saved in {output_dir}")

def save_results(pairs, pt_diffs, ct_diffs, milp_selected_diffs, output_dir):
    """Save plaintext, ciphertext, and differential results."""
    
    # Save plaintext and ciphertext pairs to file
    pairs_filename = f"{output_dir}/plaintext_ciphertext_pairs.txt"
    save_to_file(pairs, pairs_filename)

    # Save plaintext and ciphertext differences to file
    pt_diffs_filename = f"{output_dir}/plaintext_diffs.txt"
    ct_diffs_filename = f"{output_dir}/ciphertext_diffs.txt"
    
    save_to_file(pt_diffs.items(), pt_diffs_filename)
    save_to_file(ct_diffs.items(), ct_diffs_filename)
    
    # Save selected differential paths to file
    milp_diffs_filename = f"{output_dir}/milp_selected_diffs.txt"
    save_to_file(milp_selected_diffs, milp_diffs_filename)
    
    print(f"Results saved in {output_dir}")

# ========== Main Experiment ==========

if __name__ == "__main__":
    # Parameters
    n_pairs = 1000
    rounds = 5
    key = 0x1234
    target_diff = 0x5678  # Example difference
    sbox = [0x6, 0x4, 0xC, 0x5, 0x0, 0x7, 0x2, 0xE, 0x1, 0xF, 0x3, 0xD, 0x8, 0xA, 0x9, 0xB]  # Example S-box
    output_dir = "/home/zcw/gift"
    
    # Generate plaintext-ciphertext pairs
    pairs, plaintexts = generate_plaintext_ciphertext_pairs(n_pairs, key, rounds)
    
    # Calculate differences
    pt_diffs, ct_diffs = calculate_differences(pairs)
    ddt = generate_ddt(sbox)
    
    # Traditional attack
    start_time = time.time()
    traditional_keys = traditional_attack(pairs, ddt, target_diff)
    traditional_time = time.time() - start_time
    
    # MILP optimized attack
    start_time = time.time()
    milp_selected_diffs = milp_optimized_attack(pairs, ddt, target_diff)
    improved_time = time.time() - start_time
    
    # Success rates and probabilities
    success_rate_traditional = len(traditional_keys) / 0xFFFF * 100
    success_rate_improved = len(milp_selected_diffs) / 0xFFFF * 100
    prob_traditional = calculate_path_probability([(0x1, 0x4), (0x4, 0x6)], ddt)
    prob_improved = calculate_path_probability([(0x1, 0x5), (0x5, 0x7)], ddt)
    
    # Visualize and save all results
    visualize_results_and_save(
        [success_rate_traditional, success_rate_improved],
        [prob_traditional, prob_improved],
        [traditional_time, improved_time],
        plaintexts,
        pairs,
        pt_diffs,
        ct_diffs,
        output_dir
    )
    
    # Save results to text files
    save_results(pairs, pt_diffs, ct_diffs, milp_selected_diffs, output_dir)

